import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import os
from app.config import Config


class FaultDiagnosisModel(nn.Module):
    def __init__(self, input_size=10, hidden_size=64, num_classes=4):
        super(FaultDiagnosisModel, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_size, num_classes)
        )

        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = optim.Adam(self.parameters(), lr=0.001)

    def forward(self, x):
        return self.network(x)

    def save_model(self, filename='model.pth'):
        """保存模型"""
        path = os.path.join(Config.MODEL_FOLDER, filename)
        torch.save({
            'model_state_dict': self.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
        }, path)

    def load_model(self, filename='model.pth'):
        """加载模型"""
        path = os.path.join(Config.MODEL_FOLDER, filename)
        if os.path.exists(path):
            checkpoint = torch.load(path)
            self.load_state_dict(checkpoint['model_state_dict'])
            self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            return True
        return False

    def train_model(self, train_loader, epochs=10):
        """训练模型"""
        self.train()
        training_history = []

        for epoch in range(epochs):
            epoch_loss = 0
            correct = 0
            total = 0

            for inputs, labels in train_loader:
                self.optimizer.zero_grad()
                outputs = self(inputs)
                loss = self.criterion(outputs, labels)
                loss.backward()
                self.optimizer.step()

                epoch_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

            accuracy = 100 * correct / total
            avg_loss = epoch_loss / len(train_loader)
            training_history.append({
                'epoch': epoch + 1,
                'loss': avg_loss,
                'accuracy': accuracy
            })

        return training_history

    def evaluate(self, test_loader):
        """评估模型"""
        self.eval()
        correct = 0
        total = 0
        predictions = []
        true_labels = []

        with torch.no_grad():
            for inputs, labels in test_loader:
                outputs = self(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

                predictions.extend(predicted.numpy())
                true_labels.extend(labels.numpy())

        accuracy = 100 * correct / total
        return accuracy, predictions, true_labels
